import os
import sys
import time
import argparse
from selenium.webdriver import Chrome
from selenium.webdriver import ChromeOptions
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as ec
from selenium.webdriver.common.by import By


# Constant
BrowserDataDir = '../browserdata'
DownloadDir = '../download'
LogDir = '../log'
DriverExePath = '../driver/chromedriver.exe'
DriverLogPath = '../log/driver.log'

ImageXpath = '/html/body/div[3]/div/div[1]/div/img'
ImageLocator = (By.XPATH, ImageXpath)
ImageLocatedEc = ec.presence_of_element_located(ImageLocator)


def dir_init():
    if not os.path.exists(BrowserDataDir):
        os.mkdir(BrowserDataDir)
    if not os.path.exists(DownloadDir):
        os.mkdir(DownloadDir)
    if not os.path.exists(LogDir):
        os.mkdir(LogDir)
    return


def driver_init():
    opt = ChromeOptions()
    opt.add_argument(r'user-data-dir={}'.format(os.path.abspath(BrowserDataDir)))
    chrome = Chrome(executable_path=DriverExePath, options=opt, service_log_path='{}/log.txt'.format(LogDir))
    return chrome


def download_artwork(driver, artwork_url):
    # open page
    script = 'window.open("{}")'.format(artwork_url)
    driver.execute_script(script)

    # switch window
    chrome.switch_to.window(chrome.window_handles[-1])

    # skip if animation
    try:
        time.sleep(3)
        btn_pause = chrome.find_element_by_xpath('//*[@id="root"]/div[2]/div[2]/div/div[1]/main/section/div[1]/div/figure/div/div/button')
        btn_pause_disabled = btn_pause.get_attribute('disabled')
        if not btn_pause_disabled:
            print('this is an animation, not support to download yet')
            chrome.close()
            chrome.switch_to.window(chrome.window_handles[-1])
            return
    except Exception as e:
        pass
    finally:
        print('this is not animation')

    # find image element
    try:
        ActionChains(chrome).send_keys('v').perform()
        time.sleep(3)
        chrome.find_element_by_xpath(ImageXpath)
        print('single image')
    except Exception as e:
        ActionChains(chrome).send_keys('v').perform()
        time.sleep(3)
        print('multi image')

    # start download
    while True:
        try:
            WebDriverWait(chrome, 10, 1).until(ImageLocatedEc)
            image_element = chrome.find_element_by_xpath(ImageXpath)
            image_url = image_element.get_attribute('src')
            print(image_url)
            ActionChains(chrome).send_keys(Keys.DOWN).perform()
            download_artwork_image(chrome, image_url)
        except:
            break

    # exit
    print('artwork download complete')
    chrome.close()
    chrome.switch_to.window(chrome.window_handles[-1])
    return


def download_artwork_image(driver, image_url):
    # open page
    script = 'window.open("{}")'.format(image_url)
    driver.execute_script(script)

    # switch window
    chrome.switch_to.window(chrome.window_handles[-1])

    # screenshot
    title = chrome.title
    image_name = title.split('.')[0]
    image_element = chrome.find_element_by_xpath('/html/body/img')
    image_element.screenshot('{}/{}.png'.format(DownloadDir, image_name))

    # exit
    chrome.close()
    chrome.switch_to.window(chrome.window_handles[-1])
    return


if __name__ == '__main__':

    dir_init()
    chrome = driver_init()
    chrome.get('https://www.pixiv.net')
    download_artwork(chrome, 'https://www.pixiv.net/artworks/85485335')

